{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7918c0c5",
   "metadata": {},
   "source": [
    "# Structured outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6983408b",
   "metadata": {},
   "source": [
    "## with_structured_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1358e622",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Optional\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "class UserProfile(BaseModel):\n",
    "    name: str = Field(\n",
    "        description=\"The user's preferred name.\",\n",
    "    )\n",
    "    profession: str = Field(\n",
    "        description=\"The user's profession or job title.\",\n",
    "    )\n",
    "    seniority: Optional[str] = Field(\n",
    "        description=\"The user's seniority level, e.g., 'mid-level', 'senior', etc. If unknown, guess it based on profession description\",\n",
    "    )\n",
    "    languages: List[str] = Field(\n",
    "        description=\"A list of programming languages the user is proficient in.\",\n",
    "    )\n",
    "    frameworks: List[str] = Field(\n",
    "        description=\"A list of frameworks or major technologies the user uses.\",\n",
    "    )\n",
    "    current_project: Optional[str] = Field(\n",
    "        description=\"The user's current main project.\",\n",
    "    )\n",
    "    skills: List[str] = Field(\n",
    "        description=\"A list of the user's technical or professional skills.\",\n",
    "    )\n",
    "    current_interest: Optional[str] = Field(\n",
    "        description=\"The user's current main interest or focus, e.g., a topic or technology they are actively exploring or asking about.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "236774bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "structured_model = model.with_structured_output(UserProfile)\n",
    "\n",
    "profile = structured_model.invoke([HumanMessage(content=\"\"\"\n",
    "Hey, how are you? I am Evgeny, I am a software engineer and I need your help with \n",
    "my project. This is a TaskManager Java Spring Boot application. I do not get how\n",
    "to configure security there!\n",
    "\"\"\")])\n",
    "print(profile.model_dump_json(indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7486fa0",
   "metadata": {},
   "source": [
    "{\n",
    "  \"$schema\": \"https://json-schema.org/draft/2020-12/schema\",\n",
    "  \"type\": \"object\",\n",
    "  \"properties\": {\n",
    "    \"name\": {\n",
    "      \"type\": \"string\",\n",
    "      \"description\": \"The user's preferred name.\"\n",
    "    },\n",
    "    \"profession\": {\n",
    "      \"type\": \"string\",\n",
    "      \"description\": \"The user's profession or job title.\"\n",
    "    },\n",
    "    \"seniority\": {\n",
    "      \"type\": [\"string\", \"null\"],\n",
    "      \"description\": \"The user's seniority level, e.g., 'mid-level', 'senior', etc. \n",
    "                         If unknown, guess it based on profession description\"\n",
    "    },\n",
    "    ...\n",
    "  },\n",
    "  \"required\": [\"name\", \"profession\", \"languages\", \"frameworks\", \"skills\"],\n",
    "  \"additionalProperties\": false\n",
    "}\n",
    "\n",
    "Hey, how are you? I am Evgeny, I am a software engineer and I need your help with \n",
    "my project. This is a TaskManager Java Spring Boot application. I do not get how\n",
    "to configure security there!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abf26a2e",
   "metadata": {},
   "source": [
    "### Why Classical Prompt-Based Structured Output Is Problematic\n",
    "\n",
    "1. No Structural Guarantees\n",
    "\n",
    "The LLM still generates raw text. Parsing is a best-effort guess — not a guarantee. What can go wrong:\n",
    "\n",
    "- Missing fields\n",
    "- Wrong types (\"25\" instead of 25)\n",
    "- Extra fields added\n",
    "- Malformed JSON\n",
    "- Incomplete output\n",
    "\n",
    "2. Full Regeneration on Every Prompt\n",
    "\n",
    "LLMs rebuild the entire structure every time. It can’t update just one field.\n",
    "\n",
    "Why it’s a problem:\n",
    "\n",
    "- You can’t merge or incrementally fill data\n",
    "- Schema updates break old extractions\n",
    "- You risk data loss on every re-generation\n",
    "- Price and Time - you are paying for every token that was repeated, bigger prompts tak more time to process\n",
    "\n",
    "\n",
    "3. Parsing Is Still Fragile\n",
    "\n",
    "Libraries help, but they can’t prevent:\n",
    "\n",
    "- Markdown fences (```json) around JSON\n",
    "- Trailing commas, weird spacing\n",
    "- Fields outside the expected structure"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c14a243d",
   "metadata": {},
   "source": [
    "## Tool calling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b6880d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "model_with_tools = model.bind_tools(\n",
    "    [UserProfile], \n",
    "    strict=True\n",
    ")\n",
    "\n",
    "response = model_with_tools.invoke([HumanMessage(content=\"\"\"\n",
    "Hey, how are you? I am Evgeny, I am a software engineer and I need your help with \n",
    "my project. This is a TaskManager Java Spring Boot application. I do not get how\n",
    "to configure security there!\n",
    "\"\"\")])\n",
    "print(response.model_dump_json(indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01292680",
   "metadata": {},
   "outputs": [],
   "source": [
    "response.additional_kwargs[\"tool_calls\"][0][\"function\"][\"arguments\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b703d01",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "tool_call_data_raw = response.additional_kwargs[\"tool_calls\"][0][\"function\"][\"arguments\"]\n",
    "tool_call_data = json.loads(tool_call_data_raw)\n",
    "\n",
    "profile = UserProfile(**tool_call_data)\n",
    "\n",
    "print(profile.model_dump_json(indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b026517b",
   "metadata": {},
   "source": [
    "### What about Tooling ?\n",
    "\n",
    "1. No Structural Guarantees → Improved\n",
    "\n",
    "- Tool calling enforces a valid schema: arguments must match expected structure.\n",
    "- Eliminates issues like malformed JSON, markdown wrappers, and invalid formatting.\n",
    "\n",
    "2. Parsing Is Still Fragile → Improved\n",
    "\n",
    "- Tool calling returns clean structured data.\n",
    "- No need for manual parsing, cleanup, or error-prone regex hacks.\n",
    "\n",
    "3. Full Regeneration on Every Prompt - Still a Problem\n",
    "\n",
    "- The model still regenerates the full object each time.\n",
    "- Partial updates and field-level diffs are not supported.\n",
    "- You still pay for repeated content and risk data loss if the model forgets a field."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5d49929",
   "metadata": {},
   "source": [
    "## update profile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "649848a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trustcall import create_extractor\n",
    "from langchain_core.messages import HumanMessage\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "trustcall_extractor = create_extractor(model, tools=[UserProfile])\n",
    "\n",
    "\n",
    "# Invoke the extractor\n",
    "result = trustcall_extractor.invoke([HumanMessage(content=\"\"\"\n",
    "Hey, how are you? I am Evgeny, I am a software engineer and I need your help with \n",
    "my project. This is a TaskManager Java Spring Boot application. I do not get how\n",
    "to configure security there!\n",
    "\"\"\")])\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a920a45e",
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = result[\"responses\"]\n",
    "schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c344c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile = schema[0].model_dump()\n",
    "profile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f4f60e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "message = \"\"\"\n",
    "Update the memory (JSON doc) to incorporate new information from the following conversation:\n",
    "<conversation>\n",
    "I also can develop in Kotlin and ReactJS\n",
    "</conversation>\n",
    "\"\"\"\n",
    "\n",
    "result = trustcall_extractor.invoke(\n",
    "    {\n",
    "        \"messages\": [HumanMessage(content=message)], \n",
    "        \"existing\": {\"UserProfile\": profile}\n",
    "    }\n",
    ")  \n",
    "\n",
    "profile_after_update = result[\"responses\"][0].model_dump()\n",
    "profile_after_update"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fed36969",
   "metadata": {},
   "source": [
    "## Chat bot with profile updates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f34f1da0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import StateGraph, MessagesState, START, END\n",
    "from langchain_core.runnables.config import RunnableConfig\n",
    "from langgraph.store.base import BaseStore\n",
    "from langchain_core.messages import SystemMessage\n",
    "from langgraph.store.memory import InMemoryStore\n",
    "\n",
    "long_term_memory = InMemoryStore()\n",
    "\n",
    "user_id = \"1\"\n",
    "namespace = (user_id, \"memory\")\n",
    "key = \"profile\"\n",
    "\n",
    "profile = None\n",
    "\n",
    "long_term_memory.put(namespace, key, profile)\n",
    "\n",
    "\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "### Nodes\n",
    "\n",
    "def chat(state: MessagesState, config: RunnableConfig, store: BaseStore):\n",
    "    user_id = config[\"configurable\"][\"user_id\"]\n",
    "\n",
    "    # Retrieve memory from the store\n",
    "    profile = store.get((user_id, \"memory\"), \"profile\")\n",
    "\n",
    "    # Extract the actual memory content if it exists and add a prefix\n",
    "    if profile:\n",
    "        profile_content = profile.value\n",
    "    else:\n",
    "        profile_content = None\n",
    "\n",
    "    # Format the memory in the system prompt\n",
    "    system_msg = f\"\"\"\n",
    "    You are a helpful assistant with memory capabilities.\n",
    "    If user-specific memory is available, use it to personalize \n",
    "    your responses based on what you know about the user.\n",
    "    \n",
    "    Your goal is to provide relevant, friendly, and tailored \n",
    "    assistance that reflects the user’s preferences, context, and past interactions.\n",
    "\n",
    "    If the user’s name or relevant personal context is available, always personalize your responses by:\n",
    "        – Addressing the user by name (e.g., \"Sure, Bob...\") when appropriate\n",
    "        – Referencing known projects, tools, or preferences (e.g., \"your MCP  server typescript based project\")\n",
    "        – Adjusting the tone to feel friendly, natural, and directly aimed at the user\n",
    "\n",
    "    Avoid generic phrasing when personalization is possible. For example, instead of \"In TypeScript apps...\" say \"Since your project is built with TypeScript...\"\n",
    "\n",
    "    Use personalization especially in:\n",
    "        – Greetings and transitions\n",
    "        – Help or guidance tailored to tools and frameworks the user uses\n",
    "        – Follow-up messages that continue from past context\n",
    "\n",
    "    Always ensure that personalization is based only on known user details and not assumed.\n",
    "    Tailor your answers to level of seniority of the user if it's known.\n",
    "    \n",
    "    The user’s profile (which may be empty) is provided as JSON (extract all data and use it): {profile_content}\n",
    "    \"\"\"\n",
    "\n",
    "    response = model.invoke([SystemMessage(content=system_msg)] + state[\"messages\"])\n",
    "\n",
    "    return {\"messages\": response}\n",
    "\n",
    "\n",
    "def update_memory(state: MessagesState, config: RunnableConfig, store: BaseStore):\n",
    "    user_id = config[\"configurable\"][\"user_id\"]\n",
    "    namespace = (user_id, \"memory\")\n",
    "    key = \"profile\"\n",
    "    profile = store.get(namespace, key)\n",
    "        \n",
    "\n",
    "    if profile:\n",
    "        profile_content = profile.value\n",
    "    else:\n",
    "        profile_content = None\n",
    "\n",
    "\n",
    "    # Format the memory in the system prompt\n",
    "    system_msg = f\"\"\"\n",
    "    Update the user profile using the user's chat history.\n",
    "    Save this for future reference. If a profile already exists, just update it.\n",
    "    \"\"\"\n",
    "\n",
    "    if not profile_content:\n",
    "        print(\"Create profile\")\n",
    "        updated_profile = trustcall_extractor.invoke(state['messages'])\n",
    "    else:\n",
    "        print(\"Update profile\")\n",
    "        updated_profile = trustcall_extractor.invoke(\n",
    "            {\n",
    "                \"messages\": [SystemMessage(content=system_msg)] + state['messages'], \n",
    "                \"existing\": {\"UserProfile\": profile_content}\n",
    "            }\n",
    "        )\n",
    "\n",
    "    store.put(namespace, key, updated_profile[\"responses\"][0].model_dump())\n",
    "\n",
    "\n",
    "# Define the graph\n",
    "builder = StateGraph(MessagesState)\n",
    "builder.add_node(\"chat\", chat)\n",
    "builder.add_node(\"update_memory\", update_memory)\n",
    "builder.add_edge(START, \"chat\")\n",
    "builder.add_edge(\"chat\", \"update_memory\")\n",
    "builder.add_edge(\"update_memory\", END)\n",
    "\n",
    "short_term_memory = MemorySaver()\n",
    "\n",
    "graph = builder.compile(checkpointer=short_term_memory, store=long_term_memory)\n",
    "\n",
    "display(Image(graph.get_graph(xray=1).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60697258",
   "metadata": {},
   "outputs": [],
   "source": [
    "# start a new conversation\n",
    "thread = {\"configurable\": {\"thread_id\": \"1\", \"user_id\": \"1\"}}\n",
    "\n",
    "# define intiial user request\n",
    "initial_input = {\"messages\": HumanMessage(content=\"\"\"\n",
    "Hey, I am Evgeny, a junior software engineer and I need your help with my project. \n",
    "This is a TaskManager Java Spring Boot application.\n",
    "\"\"\")}\n",
    "\n",
    "# run the graph and stream in values mode\n",
    "for event in graph.stream(initial_input, thread, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "839ae929",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = thread[\"configurable\"][\"user_id\"]\n",
    "namespace = (user_id, \"memory\")\n",
    "key = \"profile\"\n",
    "profile = long_term_memory.get(namespace, key)\n",
    "profile.dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7df72908",
   "metadata": {},
   "outputs": [],
   "source": [
    "# start a new conversation\n",
    "thread = {\"configurable\": {\"thread_id\": \"2\", \"user_id\": \"1\"}}\n",
    "\n",
    "# define intiial user request\n",
    "initial_input = {\"messages\": HumanMessage(content=\"\"\"\n",
    "I also use Kotlin and I do not get how to set up security configuration in my project!\n",
    "\"\"\")}\n",
    "\n",
    "# run the graph and stream in values mode\n",
    "for event in graph.stream(initial_input, thread, stream_mode=\"values\"):\n",
    "    event['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1618b00",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = thread[\"configurable\"][\"user_id\"]\n",
    "namespace = (user_id, \"memory\")\n",
    "key = \"profile\"\n",
    "profile = long_term_memory.get(namespace, key)\n",
    "profile.dict()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
